#pragma once

#include "Net.h"
#include "Message.h"
#include <mutex>
#include <unordered_map>

namespace rpc
{

    class Callback
    {
    public:
        using ptr = std::shared_ptr<Callback>;

        virtual void onMessage(const BaseConnection::ptr &conn, BaseMessage::ptr &msg) = 0;
    };

    template <typename T>
    class CallbackT : public Callback
    {
    public:
        using ptr = std::shared_ptr<CallbackT<T>>;
        using MessageCallback = std::function<void(const BaseConnection::ptr &, std::shared_ptr<T> &)>;

        CallbackT(const MessageCallback &cb) : _handle(cb) {}

        virtual void onMessage(const BaseConnection::ptr &conn, BaseMessage::ptr &msg) override
        {
            auto typeMsg = std::dynamic_pointer_cast<T>(msg);
            _handle(conn, typeMsg);
        }

    private:
        MessageCallback _handle;
    };

    class Dispatcher
    {
    public:
        using ptr = std::shared_ptr<Dispatcher>;

        template <typename T>
        void registerHandler(MType mtype, const typename CallbackT<T>::MessageCallback &handler)
        {
            std::unique_lock<std::mutex> lock(_mtx);
            _handlerMap.insert(std::make_pair(mtype, std::make_shared<CallbackT<T>>(handler)));
        }

        void onMessage(const BaseConnection::ptr &conn, BaseMessage::ptr &msg)
        {
            std::unique_lock<std::mutex> lock(_mtx);
            auto it = _handlerMap.find(msg->mtype());
            if (it != _handlerMap.end())
            {
                return it->second->onMessage(conn, msg);
            }

            ELOG("未知消息类型")
            conn->shutdown();
        }

    private:
        std::mutex _mtx;
        std::unordered_map<MType, Callback::ptr> _handlerMap;
    };

} // rpc